from python_a2a.server import A2AServer
from python_a2a.mcp import FastMCPAgent
from python_a2a.models import Message, TextContent, MessageRole
from python_a2a import run_server
import re
import asyncio
import logging

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# YFinance Agent for stock prices
class YFinanceAgent(A2AServer, FastMCPAgent):
    """Agent that provides stock price information."""
    
    def __init__(self):
        super().__init__(url="http://localhost:5004")
        FastMCPAgent.__init__(
            self,
            mcp_servers={"finance": "http://localhost:5002"}
        )
    
    def handle_message(self, message):
        """同步处理消息"""
        logger.info(f"处理YFinance消息: {message}")
        
        try:
            if hasattr(message.content, 'text'):
                text = message.content.text
                
                # 提取股票代码
                ticker = self._extract_ticker(text)
                
                if not ticker:
                    return Message(
                        content=TextContent(text="请提供股票代码来查询价格。"),
                        role=MessageRole.AGENT,
                        parent_message_id=getattr(message, 'message_id', None),
                        conversation_id=getattr(message, 'conversation_id', None)
                    )
                
                # 同步调用MCP工具获取价格
                try:
                    # 由于MCP工具可能是异步的，我们需要在同步环境中运行异步代码
                    price_info = asyncio.run(self._get_stock_price_async(ticker))
                    
                    if "error" in price_info:
                        return Message(
                            content=TextContent(text=f"获取 {ticker} 价格时出错: {price_info['error']}"),
                            role=MessageRole.AGENT,
                            parent_message_id=getattr(message, 'message_id', None),
                            conversation_id=getattr(message, 'conversation_id', None)
                        )
                    
                    return Message(
                        content=TextContent(
                            text=f"{ticker} 当前交易价格为 {price_info['price']:.2f} {price_info['currency']}。"
                        ),
                        role=MessageRole.AGENT,
                        parent_message_id=getattr(message, 'message_id', None),
                        conversation_id=getattr(message, 'conversation_id', None)
                    )
                    
                except Exception as e:
                    logger.error(f"调用MCP工具时出错: {e}")
                    # 返回模拟数据作为备用
                    return Message(
                        content=TextContent(text=f"{ticker} 当前交易价格约为 150.00 USD (模拟数据)"),
                        role=MessageRole.AGENT,
                        parent_message_id=getattr(message, 'message_id', None),
                        conversation_id=getattr(message, 'conversation_id', None)
                    )
        
        except Exception as e:
            logger.error(f"处理消息时出错: {e}")
            return Message(
                content=TextContent(text="处理您的请求时出现了错误。"),
                role=MessageRole.AGENT,
                parent_message_id=getattr(message, 'message_id', None),
                conversation_id=getattr(message, 'conversation_id', None)
            )
        
        # 默认响应
        return Message(
            content=TextContent(text="我可以提供股票代码的价格信息。"),
            role=MessageRole.AGENT,
            parent_message_id=getattr(message, 'message_id', None),
            conversation_id=getattr(message, 'conversation_id', None)
        )
    
    async def _get_stock_price_async(self, ticker: str):
        """异步获取股票价格"""
        try:
            # 调用MCP工具获取价格
            price_info = await self.call_mcp_tool("finance", "get_stock_price", ticker=ticker)
            return price_info
        except Exception as e:
            logger.error(f"MCP工具调用失败: {e}")
            # 返回模拟数据
            return {
                "price": 150.00,
                "currency": "USD",
                "ticker": ticker
            }
    
    def _extract_ticker(self, text: str) -> str:
        """从文本中提取股票代码"""
        # 匹配大写字母组成的股票代码 (1-5个字母)
        ticker_match = re.search(r'\b([A-Z]{1,5})\b', text)
        if ticker_match:
            return ticker_match.group(1)
        
        # 尝试从常见模式中提取
        patterns = [
            r"([A-Z]{1,5})的当前价格",
            r"([A-Z]{1,5})的价格",
            r"price of ([A-Z]{1,5})",
            r"price for ([A-Z]{1,5})"
        ]
        
        for pattern in patterns:
            match = re.search(pattern, text, re.IGNORECASE)
            if match:
                return match.group(1)
        
        return ""

if __name__ == "__main__":
    # 在5004端口运行此代理
    logger.info("YFinance股票价格代理正在 http://localhost:5004 上运行")
    run_server(YFinanceAgent(), port=5004, debug=True)